-- 使用贪婪算法解决集合覆盖问题
--   贪婪算法能大致解决问题，实现起来容易，得到的结果与正确结果相当接近

local Set = {}
Set.__index = Set

function Set:add(val)
    if type(assert(val)) == "table" then
        for _, v in ipairs(val) do
            self.keys[v] = true
            self.count = self.count + 1
        end
    else
        self.keys[val] = true
        self.count = self.count + 1
    end
end

function Set:del(val)
    if type(assert(val)) == "table" then
        for _, v in ipairs(val) do
            self.keys[v] = nil
            self.count = self.count - 1
        end
    else
        self.keys[val] = nil
        self.count = self.count - 1
    end
end

function Set:mem(val)
    return self.keys[val] and true or false
end

function Set:all_keys()
    return self.keys
end

function Set:len()
    return self.count
end

function Set:init()
    self.count = 0
    self.keys = {}
end

function Set.new(...)
    local set = setmetatable({}, Set)
    set:init()
    if select("#", ...) > 0 then
        set:add({...})
    end
    return set
end

function Set.union(s1, s2)
    local set = Set.new()
    for k in pairs(s1:all_keys()) do
        set:add(k)
    end
    for k in pairs(s2:all_keys()) do
        set:add(k)
    end
    return set
end

function Set.intersection(s1, s2)
    local set = Set.new()
    for k2 in pairs(s2:all_keys()) do
        if s1:mem(k2) then
            set:add(k2)
        end
    end
    return set
end

function Set.difference(s1, s2)
    local set = Set.new()
    for k1 in pairs(s1:all_keys()) do
        set:add(k1)
    end
    for k2 in pairs(s2:all_keys()) do
        if set:mem(k2) then
            set:del(k2)
        end
    end
    return set
end

function Set.tostring(s)
    local arr = {}
    for k in pairs(s:all_keys()) do
        arr[#arr+1] = k
    end
    return table.concat(arr, ",")
end

Set.__add = Set.union
Set.__mul = Set.intersection
Set.__sub = Set.difference
Set.__tostring = Set.tostring

-- 需要被覆盖的州
local states_needed = Set.new("mt", "wa", "or", "id", "nv", "ut", "ca", "az")
-- 各广播台覆盖的州
local stations = {}
stations["kone"] = Set.new("id", "nv", "ut")
stations["ktwo"] = Set.new("wa", "id", "mt")
stations["kthree"] = Set.new("or", "nv", "ca")
stations["kfour"] = Set.new("nv", "ut")
stations["kfive"] = Set.new("ca", "az")

local function _greedy_algorithms_calc_station()
    local final_stations = Set.new()

    local needed = states_needed
    while needed:len() > 0 do
        local states_covered = Set.new()
        local best_station
        -- 遍历所有集合从中选择覆盖了最多的未覆盖州的广播台
        for station, states in pairs(stations) do
            local covered = states * needed
            if covered:len() > states_covered:len() then
                best_station = station
                states_covered = covered
            end
        end

        assert(best_station)
        needed = needed - states_covered
        final_stations:add(best_station)
    end

    return final_stations
end

local function test()
    local final_stations = _greedy_algorithms_calc_station()
    -- Todo: 不明白为什么运行结果是变化的
    -- print(final_stations)

    local needed = states_needed
    local real = states_needed:len()
    local total = 0
    for station in pairs(final_stations:all_keys()) do
        needed = needed - stations[station]
        total = total + stations[station]:len()
    end
    assert(needed:len() == 0)
    print("ch08 set covering:", total/real)

    for state in pairs(states_needed:all_keys()) do
        local find
        for station in pairs(final_stations:all_keys()) do
            if stations[station]:mem(state) then
                find = true
                break
            end
        end
        assert(find, state)
    end
end

test()

